#ifndef HEAT_HPP
#define HEAT_HPP

#include "MOL.hpp"

using namespace std;
using namespace Eigen;

double f1(const double &x)
{
    if (0.45 <= x && x < 0.5)
        return 20.0 * (x - 9.0 / 20);
    else if (0.5 <= x && x < 11.0 / 20)
        return -20.0 * (x - 11.0 / 20);
    else
        return 0;
}

class heat_equ
{
private:
    double h = 0.05;
    double k, r, nu = 1;
    int m; // nk(t) ih(x)
    VectorXd U;
    vector<VectorXd> solution;

public:
    heat_equ(double _r)
    {
        r = _r;
        k = r * h * h / nu;
        m = 1.0 / h;
        U = VectorXd::Zero(m + 1);
        for (int i = 0; i <= m; i++)
            U(i) = f1(i * h);
        solution.push_back(U);
    }
    void plot(int index)
    {
        U = solution[index];
        ofstream outfile;
        outfile.open("output.m");
        outfile << "x = 0:" << h << ":1;" << endl;
        outfile << "y = [";
        for (int i = 0; i <= m; i++)
        {
            outfile << U(i) << ",";
        }
        outfile << "];" << endl;
        outfile << "plot(x, y,\".-\");" << endl;
        outfile.close();
    }

    void theta_method(double theta)
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i, i) = -2;
            A(i + 1, i) = 1;
            A(i, i + 1) = 1;
        }
        A(m, m) = -2;

        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
        MatrixXd B = I - theta * r * A;
        MatrixXd C = I + (1 - theta) * r * A;
        int n = 0;
        while (n < 15)
        {
            U = B.inverse() * C * U;
            solution.push_back(U);
            n++;
        }
    }

    void FTCS()
    {
        theta_method(0);
    }

    void crank_nicolson()
    {
        theta_method(0.5);
    }

    void BTCS()
    {
        theta_method(1);
    }

    void gauss_legendre()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i, i) = -2;
            A(i + 1, i) = 1;
            A(i, i + 1) = 1;
        }
        A(m, m) = -2;
        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);

        int n = 0;
        while (n < 15)
        {
            VectorXd y1 = (I - 0.5 * r * A).lu().solve(r * A * U);
            U = U + y1;
            solution.push_back(U);
            n++;
        }
    }

    void collocation()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i, i) = -2;
            A(i + 1, i) = 1;
            A(i, i + 1) = 1;
        }
        A(m, m) = -2;
        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);

        int n = 0;
        while (n < 15)
        {
            // VectorXd y1 = ((I - 1.0 / 3 * r * A) * A.inverse() * (I - 5.0 / 12 * r * A) + r / 12.0 * (I + 1.0 / 3 * r * A)).lu().solve(r * U);
            VectorXd y1 = (I - 5.0 / 12 * r * A + r / 12.0 * A * (I - 1.0 / 3 * r * A).inverse() * (I + 1.0 / 3 * r * A)).lu().solve(r * A * U);
            VectorXd y2 = (I - 1.0 / 3 * r * A).inverse() * (I + 1.0 / 3 * r * A) * y1;
            U = U + (0.75 * y1 + 0.25 * y2);
            solution.push_back(U);
            n++;
        }
    }
};

#endif